// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/core/plan_node.h>
#include <pollux/exec/operator.h>

namespace kumo::pollux::exec {
    /// Implements a merge operator to aggregate the metadata outputs from multiple
    /// table write operators and produces the aggregated result.
    class TableWriteMerge : public Operator {
    public:
        TableWriteMerge(
            int32_t operatorId,
            DriverCtx *driverCtx,
            const std::shared_ptr<const core::TableWriteMergeNode> &
            tableWriteMergeNode);

        void initialize() override;

        BlockingReason isBlocked(ContinueFuture * /* future */) override {
            return BlockingReason::kNotBlocked;
        }

        void addInput(RowVectorPtr input) override;

        void noMoreInput() override;

        virtual bool needsInput() const override {
            return true;
        }

        RowVectorPtr getOutput() override;

        bool isFinished() override {
            return finished_;
        }

    private:
        // Creates non-last output with fragments and last commit context only.
        RowVectorPtr createFragmentsOutput();

        // Creates json encoded string of last commit context with specified
        // 'lastOutput' flag.
        std::string createTableCommitContext(bool lastOutput) const;

        // Creates the last output and fragment columns must be null.
        RowVectorPtr createLastOutput();

        // Check if the input is statistics input.
        bool isStatistics(RowVectorPtr input);

        std::unique_ptr<Operator> aggregation_;
        bool finished_{false};
        // The sum of written rows.
        int64_t numRows_{0};
        std::queue<VectorPtr> fragmentVectors_;
        melon::Dynamic lastCommitContext_;
    };
} // namespace kumo::pollux::exec
